/*****************************************************************************
 * DotNetProfiler
 * 
 * Copyright (c) 2006 Scott Hackett
 * 
 * This software is provided 'as-is', without any express or implied warranty.
 * In no event will the author be held liable for any damages arising from the
 * use of this software. Permission to use, copy, modify, distribute and sell
 * this software for any purpose is hereby granted without fee, provided that
 * the above copyright notice appear in all copies and that both that
 * copyright notice and this permission notice appear in supporting
 * documentation.
 * 
 * Scott Hackett (code@scotthackett.com)
 *****************************************************************************/

#include "stdafx.h"
#include <assert.h>
#include "winnt.h"
#include "Profiler.h"

#pragma warning (disable: 4996) 

// --> Static locals <---------------------------------------------------------

static CProfiler* gICorProfilerCallback = NULL;

// --> Implementation <--------------------------------------------------------

CProfiler::CProfiler() 
{
    InitializeCriticalSection(&mLock);
	mSampleFile = NULL;
    mMethodFile = NULL;
    mNumSamples = 0;
    mNumMethods = 0;
    mTimeOffset = __rdtsc();
}

HRESULT CProfiler::FinalConstruct()
{
    EnterCriticalSection(&mLock);
    char outputName[512];
    if (GetEnvironmentVariableA("PROFILER_OUTPUT", outputName, sizeof(outputName)) != 0)
    {
        char name[512];
        strncpy(name, outputName, sizeof(name));
        strncat(name, "_samples.dat", sizeof(name));
        mSampleFile = fopen(name, "wb");

        strncpy(name, outputName, sizeof(name));
        strncat(name, "_methods.dat", sizeof(name));
        mMethodFile = fopen(name, "wb");
    }
    LeaveCriticalSection(&mLock);
	return S_OK;
}

void CProfiler::FinalRelease()
{
    EnterCriticalSection(&mLock);
    if (mNumSamples > 0)
    {
        fwrite(mSamples, sizeof(Sample), mNumSamples, mSampleFile);
    }
	fclose(mSampleFile);

    if (mNumMethods > 0)
    {
        fwrite(mMethods, sizeof(Method), mNumMethods, mMethodFile);
    }
	fclose(mMethodFile);
    LeaveCriticalSection(&mLock);
}

// -->  Global callbacks <-----------------------------------------------------

void __stdcall FunctionEnterGlobal(FunctionID functionID, UINT_PTR clientData, COR_PRF_FRAME_INFO frameInfo, COR_PRF_FUNCTION_ARGUMENT_INFO *argInfo)
{
    if (gICorProfilerCallback != NULL)
	{
        gICorProfilerCallback->Enter(functionID, clientData, frameInfo, argInfo);
	}
}

void __stdcall FunctionLeaveGlobal(FunctionID functionID, UINT_PTR clientData, COR_PRF_FRAME_INFO frameInfo, COR_PRF_FUNCTION_ARGUMENT_RANGE *retvalRange)
{
    if (gICorProfilerCallback != NULL)
	{
        gICorProfilerCallback->Leave(functionID, clientData, frameInfo, retvalRange);
	}
}

void __stdcall FunctionTailcallGlobal(FunctionID functionID, UINT_PTR clientData, COR_PRF_FRAME_INFO frameInfo)
{
    if (gICorProfilerCallback != NULL)
	{
        gICorProfilerCallback->Tailcall(functionID, clientData, frameInfo);
	}
}

// --> Naked methods <---------------------------------------------------------

void _declspec(naked) FunctionEnterNaked(FunctionID functionID, UINT_PTR clientData, COR_PRF_FRAME_INFO func, COR_PRF_FUNCTION_ARGUMENT_INFO *argumentInfo)
{
    __asm
    {
        push    ebp                 // Create a standard frame
        mov     ebp,esp
        pushad                      // Preserve all registers

        mov     eax,[ebp+0x14]      // argumentInfo
        push    eax
        mov     ecx,[ebp+0x10]      // func
        push    ecx
        mov     edx,[ebp+0x0C]      // clientData
        push    edx
        mov     eax,[ebp+0x08]      // functionID
        push    eax
        call    FunctionEnterGlobal

        popad                       // Restore all registers
        pop     ebp                 // Restore EBP
        ret     16
    }
}

void _declspec(naked) FunctionLeaveNaked(FunctionID functionID, UINT_PTR clientData, COR_PRF_FRAME_INFO func, COR_PRF_FUNCTION_ARGUMENT_RANGE *retvalRange)
{
    __asm
    {
        push    ebp                 // Create a standard frame
        mov     ebp,esp
        pushad                      // Preserve all registers

        mov     eax,[ebp+0x14]      // argumentInfo
        push    eax
        mov     ecx,[ebp+0x10]      // func
        push    ecx
        mov     edx,[ebp+0x0C]      // clientData
        push    edx
        mov     eax,[ebp+0x08]      // functionID
        push    eax
        call    FunctionLeaveGlobal

        popad                       // Restore all registers
        pop     ebp                 // Restore EBP
        ret     16
    }
}

void _declspec(naked) FunctionTailcallNaked(FunctionID functionID, UINT_PTR clientData, COR_PRF_FRAME_INFO func)
{
    __asm
    {
        push    ebp                 // Create a standard frame
        mov     ebp,esp
        pushad                      // Preserve all registers

        mov     eax,[ebp+0x14]      // argumentInfo
        push    eax
        mov     ecx,[ebp+0x10]      // func
        push    ecx
        mov     edx,[ebp+0x0C]      // clientData
        push    edx
        mov     eax,[ebp+0x08]      // functionID
        push    eax
        call    FunctionTailcallGlobal

        popad                       // Restore all registers
        pop     ebp                 // Restore EBP
        ret     16
    }
}

// --> Function mapper <-------------------------------------------------------

// this function is called by the CLR when a function has been mapped to an ID
UINT_PTR CProfiler::FunctionMapper(FunctionID functionID, BOOL *pbHookFunction)
{
	// make sure the global reference to our profiler is valid.  Forward this
	// call to our profiler object
    if (gICorProfilerCallback != NULL)
	{
        gICorProfilerCallback->MapFunction(functionID);
	}

	// we must return the function ID passed as a parameter
	return (UINT_PTR)functionID;
}

// the static function called by .Net when a function has been mapped to an ID
void CProfiler::MapFunction(FunctionID functionID)
{
	// see if this function is in the map
    EnterCriticalSection(&mLock);
	std::set<FunctionID>::iterator iter = mFunctionMap.find(functionID);
	if (iter == mFunctionMap.end())
	{
		// declared in this block so they are not created if the function is found
		WCHAR szMethodName[NAME_BUFFER_SIZE];
		const WCHAR* p = NULL;
		USES_CONVERSION;

		// get the method name
		HRESULT hr = GetFullMethodName(functionID, szMethodName, NAME_BUFFER_SIZE); 
		if (FAILED(hr))
		{
			// if we couldn't get the function name, then log it
			// LogString(MSG_ERROR, 0, "Unable to find the name for %i.", functionID);
			return;
		}

		// add it to the set
		mFunctionMap.insert(functionID);
		LogMethod(functionID, W2A(szMethodName));
	}
    LeaveCriticalSection(&mLock);
}

// --> Callback methods <------------------------------------------------------

void CProfiler::Enter(FunctionID functionID, UINT_PTR clientData, COR_PRF_FRAME_INFO frameInfo, COR_PRF_FUNCTION_ARGUMENT_INFO *argumentInfo)
{
	LogSample(MSG_ENTER, functionID);
}

void CProfiler::Leave(FunctionID functionID, UINT_PTR clientData, COR_PRF_FRAME_INFO frameInfo, COR_PRF_FUNCTION_ARGUMENT_RANGE *argumentRange)
{
	LogSample(MSG_LEAVE, functionID);
}

void CProfiler::Tailcall(FunctionID functionID, UINT_PTR clientData, COR_PRF_FRAME_INFO frameInfo)
{
	LogSample(MSG_TAIL, functionID);
}

STDMETHODIMP CProfiler::ExceptionThrown(ObjectID thrownObjectID)
{
	LogSample(MSG_ExceptionThrown, thrownObjectID);
    return S_OK;
}

STDMETHODIMP CProfiler::ExceptionUnwindFunctionEnter(FunctionID functionID)
{
	LogSample(MSG_ExceptionUnwindFunctionEnter, functionID);
    return S_OK;
}

STDMETHODIMP CProfiler::ExceptionUnwindFunctionLeave()
{
	LogSample(MSG_ExceptionUnwindFunctionLeave, 0);
    return S_OK;
}

STDMETHODIMP CProfiler::UnmanagedToManagedTransition(FunctionID functionID, COR_PRF_TRANSITION_REASON reason)
{
	LogSample(MSG_UnmanagedToManagedTransition, functionID);
    return S_OK;
}

STDMETHODIMP CProfiler::ManagedToUnmanagedTransition(FunctionID functionID, COR_PRF_TRANSITION_REASON reason)
{
	LogSample(MSG_ManagedToUnmanagedTransition, functionID);
    return S_OK;
}

STDMETHODIMP CProfiler::Initialize(IUnknown *pICorProfilerInfoUnk)
{
	// set up our global access pointer
	gICorProfilerCallback = this;

	// get the ICorProfilerInfo interface
    HRESULT hr = pICorProfilerInfoUnk->QueryInterface(IID_ICorProfilerInfo, (LPVOID*)&mICorProfilerInfo);
    if (FAILED(hr))
	{
        return E_FAIL;
	}
	// determine if this object implements ICorProfilerInfo2
    hr = pICorProfilerInfoUnk->QueryInterface(IID_ICorProfilerInfo2, (LPVOID*)&mICorProfilerInfo2);
    if (FAILED(hr))
	{
		// we still want to work if this call fails, might be an older .NET version
		mICorProfilerInfo2.p = NULL;
	}

	// Indicate which events we're interested in.
	hr = SetEventMask();
    if (FAILED(hr))
	{
        //LogString(MSG_ERROR, 0, "Setting the event mask.");
	}

	// set the enter, leave and tailcall hooks
	if (mICorProfilerInfo2.p == NULL)
	{
		// note that we are casting our functions to the definitions for the callbacks
		hr = mICorProfilerInfo->SetEnterLeaveFunctionHooks((FunctionEnter*)&FunctionEnterNaked, (FunctionLeave*)&FunctionLeaveNaked, (FunctionTailcall*)&FunctionTailcallNaked);
		if (SUCCEEDED(hr))
		{
			hr = mICorProfilerInfo->SetFunctionIDMapper((FunctionIDMapper*)&FunctionMapper);
		}
	}
	else
	{
		hr = mICorProfilerInfo2->SetEnterLeaveFunctionHooks2(FunctionEnterNaked, FunctionLeaveNaked, FunctionTailcallNaked);
		if (SUCCEEDED(hr))
		{
			hr = mICorProfilerInfo2->SetFunctionIDMapper(FunctionMapper);
		}
	}
	// report our success or failure to the log file
    if (FAILED(hr))
	{
        //LogString(MSG_ERROR, 0, "Setting the enter, leave and tailcall hooks.");
	}
	else
	{
		LARGE_INTEGER freq;
		QueryPerformanceFrequency(&freq);
	}
    return S_OK;
}

STDMETHODIMP CProfiler::Shutdown()
{
	gICorProfilerCallback = NULL;
    return S_OK;
}

void CProfiler::LogSample(MessageID aMessage, FunctionID aFunctionID)
{
    EnterCriticalSection(&mLock);
    __int64 start = __rdtsc();

    // flush samples.
    if (mNumSamples > NUM_SAMPLES)
    {
        if (mSampleFile != NULL)
        {
            fwrite(mSamples, sizeof(Sample), mNumSamples, mSampleFile);
        }
        mNumSamples = 0;
    }

    // get current thread.
    ThreadID threadID = 0;
    if (mICorProfilerInfo != NULL)
    {
        mICorProfilerInfo->GetCurrentThreadID(&threadID);
    }

    // write out to the pipe.
    mSamples[mNumSamples].mMessage  = aMessage;
    mSamples[mNumSamples].mThread   = threadID;
    mSamples[mNumSamples].mFunction = aFunctionID;
    mSamples[mNumSamples].mCycle    = start - mTimeOffset;
    mNumSamples++;

    mTimeOffset += (__rdtsc() - start);
    LeaveCriticalSection(&mLock);
}

void CProfiler::LogMethod(FunctionID aFunctionID, const char* aString)
{
    EnterCriticalSection(&mLock);
    __int64 start = __rdtsc();

    mMethods[mNumMethods].mFunction = aFunctionID;
    strncpy(mMethods[mNumMethods].mText, aString, sizeof(mMethods[mNumMethods].mText));
    mNumMethods++;
    if (mNumMethods >= NUM_METHODS)
    {
        if (mMethodFile != NULL)
        {
            fwrite(mMethods, sizeof(Method), mNumMethods, mMethodFile);
        }
        mNumMethods = 0;
    }

    mTimeOffset += (__rdtsc() - start);
    LeaveCriticalSection(&mLock);
}

HRESULT CProfiler::SetEventMask()
{
	DWORD eventMask = 
		COR_PRF_MONITOR_ENTERLEAVE | 
        COR_PRF_MONITOR_EXCEPTIONS |
		COR_PRF_MONITOR_CODE_TRANSITIONS;
	return mICorProfilerInfo->SetEventMask(eventMask);
}

HRESULT CProfiler::GetFullMethodName(FunctionID functionID, LPWSTR wszMethod, int cMethod)
{
	IMetaDataImport* pIMetaDataImport = 0;
	HRESULT hr = S_OK;
	mdToken funcToken = 0;
	WCHAR szFunction[NAME_BUFFER_SIZE];
	WCHAR szClass[NAME_BUFFER_SIZE];

	// get the token for the function which we will use to get its name
	hr = mICorProfilerInfo->GetTokenAndMetaDataFromFunction(functionID, IID_IMetaDataImport, (LPUNKNOWN *) &pIMetaDataImport, &funcToken);
	if(SUCCEEDED(hr))
	{
		mdTypeDef classTypeDef;
		ULONG cchFunction;
		ULONG cchClass;

		// retrieve the function properties based on the token
		hr = pIMetaDataImport->GetMethodProps(funcToken, &classTypeDef, szFunction, NAME_BUFFER_SIZE, &cchFunction, 0, 0, 0, 0, 0);
		if (SUCCEEDED(hr))
		{
			// get the function name
			hr = pIMetaDataImport->GetTypeDefProps(classTypeDef, szClass, NAME_BUFFER_SIZE, &cchClass, 0, 0);
			if (SUCCEEDED(hr))
			{
				// create the fully qualified name
				_snwprintf_s(wszMethod,cMethod,cMethod,L"%s.%s",szClass,szFunction);
			}
		}
		// release our reference to the metadata
		pIMetaDataImport->Release();
	}
	return hr;
}


